#!/usr/bin/env python

import rospy
import actionlib
from geometry_msgs.msg import Twist, PoseStamped
from move_base_msgs.msg import MoveBaseAction, MoveBaseGoal
from std_srvs.srv import Empty
from pynput import keyboard
import threading
import math
import random
import time

class TeleopWithNavigation:
    def __init__(self):
        # 初始化 ROS 节点
        rospy.init_node('teleop_with_navigation', anonymous=True)
        
        # 创建 Twist 消息发布者，发布到 /cmd_vel 话题
        self.cmd_vel_pub = rospy.Publisher('/cmd_vel', Twist, queue_size=10)
        
        # 导航客户端
        self.move_base_client = actionlib.SimpleActionClient('move_base', MoveBaseAction)
        rospy.loginfo("等待move_base服务器...")
        
        # 设置操作模式
        self.mode = "manual"  # "manual" 或 "auto"
        self.auto_nav_running = False
        self.nav_thread = None
        
        # 初始化速度值
        self.linear_speed = 0.2  # 线速度 (m/s)
        self.angular_speed = 1.0  # 角速度 (rad/s)
        self.twist = Twist()
        
        # 按键状态
        self.forward = False
        self.backward = False
        self.left = False
        self.right = False
        
        # 显示说明
        self.display_instructions()

    def display_instructions(self):
        """显示操作说明"""
        rospy.loginfo("\n=== TurtleBot3 控制系统 ===")
        rospy.loginfo("方向控制:")
        rospy.loginfo("  ↑ : 前进")
        rospy.loginfo("  ↓ : 后退")
        rospy.loginfo("  ← : 左转")
        rospy.loginfo("  → : 右转")
        rospy.loginfo("模式切换:")
        rospy.loginfo("  M : 手动模式")
        rospy.loginfo("  A : 自动导航模式")
        rospy.loginfo("  空格键 : 急停")
        rospy.loginfo("  ESC : 退出程序")
        rospy.loginfo("当前模式: 手动")
        rospy.loginfo("==========================\n")

    def on_press(self, key):
        try:
            # 普通方向键处理
            if key == keyboard.Key.up:
                self.forward = True
            elif key == keyboard.Key.down:
                self.backward = True
            elif key == keyboard.Key.left:
                self.left = True
            elif key == keyboard.Key.right:
                self.right = True
            elif key == keyboard.Key.space:
                # 空格键 - 急停
                self.stop_robot()
                rospy.loginfo("急停!")
            elif key == keyboard.Key.esc:
                # ESC键 - 退出
                rospy.loginfo("程序退出!")
                return False
            elif hasattr(key, 'char'):
                if key.char in ['m', 'M']:
                    # 切换到手动模式
                    self.set_manual_mode()
                elif key.char in ['a', 'A']:
                    # 切换到自动导航模式
                    self.set_auto_mode()
        except AttributeError:
            pass
            
        if self.mode == "manual":
            self.update_speed()
        
        return True

    def on_release(self, key):
        try:
            # 释放键盘上的箭头键
            if key == keyboard.Key.up:
                self.forward = False
            elif key == keyboard.Key.down:
                self.backward = False
            elif key == keyboard.Key.left:
                self.left = False
            elif key == keyboard.Key.right:
                self.right = False
        except AttributeError:
            pass
            
        if self.mode == "manual":
            self.update_speed()

    def update_speed(self):
        """根据按键状态更新速度"""
        # 只在手动模式下更新速度
        if self.mode != "manual":
            return
            
        # 根据按键状态更新速度
        self.twist.linear.x = self.linear_speed if self.forward else (-self.linear_speed if self.backward else 0.0)
        self.twist.linear.y = 0.0
        self.twist.linear.z = 0.0
        self.twist.angular.x = 0.0
        self.twist.angular.y = 0.0
        self.twist.angular.z = self.angular_speed if self.left else (-self.angular_speed if self.right else 0.0)
        
        # 发布速度指令
        self.cmd_vel_pub.publish(self.twist)

    def set_manual_mode(self):
        """切换到手动模式"""
        if self.mode == "auto":
            # 如果当前是自动模式，则停止导航
            if self.auto_nav_running:
                self.stop_auto_navigation()
                
        self.mode = "manual"
        rospy.loginfo("已切换到手动模式")
        
        # 为安全起见，立即停止机器人
        self.stop_robot()

    def set_auto_mode(self):
        """切换到自动导航模式"""
        if self.mode == "manual":
            # 停止手动控制的任何动作
            self.stop_robot()
            
        self.mode = "auto"
        rospy.loginfo("已切换到自动导航模式")
        
        # 开始自动导航
        if not self.auto_nav_running:
            self.start_auto_navigation()

    def stop_robot(self):
        """发送停止命令"""
        twist = Twist()
        self.cmd_vel_pub.publish(twist)

    def start_auto_navigation(self):
        """启动自动导航线程"""
        if not self.auto_nav_running:
            self.auto_nav_running = True
            self.nav_thread = threading.Thread(target=self.auto_navigation_loop)
            self.nav_thread.daemon = True
            self.nav_thread.start()
            rospy.loginfo("开始自动导航")

    def stop_auto_navigation(self):
        """停止自动导航"""
        if self.auto_nav_running:
            self.auto_nav_running = False
            # 取消当前导航目标
            self.move_base_client.cancel_all_goals()
            rospy.loginfo("已停止自动导航")
            # 等待线程结束
            if self.nav_thread:
                self.nav_thread.join(timeout=1.0)

    def auto_navigation_loop(self):
        """自动导航的主循环"""
        try:
            # 等待move_base服务器启动
            self.move_base_client.wait_for_server(rospy.Duration(5.0))
            
            # 获取地图参数（这里需要根据你的实际情况调整）
            map_width = rospy.get_param('/map_width', 10.0)  # 默认地图宽度为10米
            map_height = rospy.get_param('/map_height', 10.0)  # 默认地图高度为10米
            
            # 设置探索区域的边界
            x_min, x_max = -map_width/2, map_width/2
            y_min, y_max = -map_height/2, map_height/2
            
            # 导航循环
            while self.auto_nav_running and not rospy.is_shutdown():
                # 随机生成下一个目标点
                x = random.uniform(x_min, x_max)
                y = random.uniform(y_min, y_max)
                
                # 设置导航目标
                goal = MoveBaseGoal()
                goal.target_pose.header.frame_id = "map"
                goal.target_pose.header.stamp = rospy.Time.now()
                goal.target_pose.pose.position.x = x
                goal.target_pose.pose.position.y = y
                goal.target_pose.pose.orientation.w = 1.0  # 朝向不太重要，设为默认值
                
                rospy.loginfo(f"导航到新目标点: x={x:.2f}, y={y:.2f}")
                
                # 发送目标
                self.move_base_client.send_goal(goal)
                
                # 等待直到达到目标、被取消或者模式切换
                while self.auto_nav_running and not rospy.is_shutdown():
                    # 检查导航状态
                    state = self.move_base_client.get_state()
                    if state == actionlib.GoalStatus.SUCCEEDED:
                        rospy.loginfo("成功到达目标点")
                        break
                    elif state == actionlib.GoalStatus.ABORTED:
                        rospy.loginfo("无法到达目标点，将重新规划路径")
                        break
                    
                    # 短暂休眠，避免CPU占用过高
                    time.sleep(0.5)
                
                # 在目标点稍作停留
                time.sleep(1.0)
                
        except Exception as e:
            rospy.logerr(f"自动导航发生错误: {e}")
            self.auto_nav_running = False

    def run(self):
        # 启动键盘监听
        with keyboard.Listener(on_press=self.on_press, on_release=self.on_release) as listener:
            rospy.loginfo("控制系统已启动，按ESC退出")
            listener.join()

if __name__ == '__main__':
    try:
        teleop = TeleopWithNavigation()
        teleop.run()
    except rospy.ROSInterruptException:
        pass

